import time
from typing import Callable

import numpy as np
from dora import DoraStatus
from dora_utils import closest_vertex, pairwise_distances, Status, BorrowLane, Action, Decision
from scipy.interpolate import splrep, splev
from scipy.spatial.transform import Rotation as R

from hybrid_astar_planner.HybridAStar import hybrid_astar_wrapper

# Hybrid AStar
HYBRID_ASTAR_STEP_SIZE = 1.0
HYBRID_ASTAR_MAX_ITERATIONS = 1000
HYBRID_ASTAR_COMPLETION_THRESHOLD = 0.5
HYBRID_ASTAR_ANGLE_COMPLETION_THRESHOLD = np.deg2rad(20)
HYBRID_ASTAR_RAD_STEP = 0.1
HYBRID_ASTAR_RAD_UPPER_RANGE = np.deg2rad(60)
HYBRID_ASTAR_RAD_LOWER_RANGE = np.deg2rad(60)
HYBRID_ASTAR_OBSTACLE_CLEARANCE = 0.5
HYBRID_ASTAR_LANE_WIDTH = 4.0
HYBRID_ASTAR_RADIUS = 1.2
HYBRID_ASTAR_CAR_LENGTH = 4.9
HYBRID_ASTAR_CAR_WIDTH = 1.8

# Obstacle
OBSTACLE_DISTANCE_THRESHOLD = 2 # m
OBSTACLE_PADDING = 0.5 # m

# Local Planner
LOCAL_PLANNER_NUM_GLOBAL_WAYPOINTS_AHEAD = 20
LOCAL_PLANNER_DEFAULT_SPEED = 5.0 # m/s
LOCAL_PLANNER_MAX_LAT_ACCEL = 2.0 # m/s^2

# Debug
DEBUG = False

def get_obstacle_list(obstacle_predictions, waypoints, padding=0.5):
    if len(obstacle_predictions) == 0 or len(waypoints) == 0:
        return np.empty((0, 4))
    obstacle_list = []

    distances = pairwise_distances(waypoints, obstacle_predictions[:, :2]).min(
        0
    )
    for distance, prediction in zip(distances, obstacle_predictions):
        # Use all prediction times as potential obstacles.
        if distance < OBSTACLE_DISTANCE_THRESHOLD:
            [x, y, _, _confidence, _label] = prediction
            obstacle_size = np.array(
                [
                    x - padding,
                    y - padding,
                    x + padding,
                    y + padding,
                ]
            )

            # Remove traffic light. TODO: Take into account traffic light.
            if _label != 9:
                obstacle_list.append(obstacle_size)

    if len(obstacle_list) == 0:
        return np.empty((0, 4))
    return np.array(obstacle_list)


class Operator:
    """
    Compute a `control` based on the position and the waypoints of the car.
    """

    def __init__(self):
        self.obstacles = []
        self.position = []
        self.speed = 0.0
        self.global_waypoints = []
        self.target_speed = LOCAL_PLANNER_DEFAULT_SPEED
        self.decision = None
        
        self._hyperparameters = {
            "step_size": HYBRID_ASTAR_STEP_SIZE,
            "max_iterations": HYBRID_ASTAR_MAX_ITERATIONS,
            "completion_threshold": HYBRID_ASTAR_COMPLETION_THRESHOLD,
            "angle_completion_threshold": HYBRID_ASTAR_ANGLE_COMPLETION_THRESHOLD,
            "rad_step": HYBRID_ASTAR_RAD_STEP,
            "rad_upper_range": HYBRID_ASTAR_RAD_UPPER_RANGE,
            "rad_lower_range": HYBRID_ASTAR_RAD_LOWER_RANGE,
            "obstacle_clearance": HYBRID_ASTAR_OBSTACLE_CLEARANCE,
            "lane_width": HYBRID_ASTAR_LANE_WIDTH,
            "radius": HYBRID_ASTAR_RADIUS,
            "car_length": HYBRID_ASTAR_CAR_LENGTH,
            "car_width": HYBRID_ASTAR_CAR_WIDTH,
        }

    def on_event(
        self,
        dora_event: dict,
        send_output: Callable[[str, bytes], None],
    ) -> DoraStatus:
        if dora_event["type"] == "INPUT":
            return self.on_input(dora_event, send_output)
        return DoraStatus.CONTINUE

    def on_input(
        self,
        dora_input: dict,
        send_output: Callable[[str, bytes], None],
    ):

        if dora_input["id"] == "position":
            self.position = np.frombuffer(dora_input["data"], np.float32)
            return DoraStatus.CONTINUE
            
        elif dora_input["id"] == "speed":
            self.speed = np.frombuffer(dora_input["data"], np.float32)
            return DoraStatus.CONTINUE

        elif dora_input["id"] == "obstacles":
            obstacles = np.frombuffer(dora_input["data"], np.float32).reshape(
                (-1, 5)
            )
            self.obstacles = obstacles
            return DoraStatus.CONTINUE
        
        elif dora_input["id"] == "decision_making":
            self.decision = Decision.from_bytes(dora_input["data"])
            return DoraStatus.CONTINUE

        elif dora_input["id"] == "global_waypoints":
            waypoints = np.frombuffer(dora_input["data"], np.float32).reshape((-1, 4))
            self.target_speed = waypoints[-1, -1] # TODO: use non-uniform speed
            self.global_waypoints = waypoints[:, :2]

            if len(self.global_waypoints) > 0 and len(self.position) > 0:
                (waypoints, target_speeds, success) = self.run(
                    time.time()
                )

                # if len(waypoints) == 0:
                #     return DoraStatus.CONTINUE
                
                local_points_x = waypoints[:, 0]
                local_points_y = waypoints[:, 1]

                # local planner waypoints
                local_points = (
                    np.concatenate([local_points_x, local_points_y, target_speeds])
                    .reshape((3, -1))
                    .T.astype(np.float32)
                )

                send_output(
                    "waypoints",
                    local_points.tobytes(),
                    dora_input["metadata"],
                )
            return DoraStatus.CONTINUE
        
        else:
            return DoraStatus.CONTINUE

    def run(self, _ttd=None):
        """Runs the planner.

        Note:
            The planner assumes that the world is up-to-date.

        Returns:
            :py:class:`~pylot.planning.waypoints.Waypoints`: Waypoints of the
            planned trajectory.
        """
        # Remove already past waypoints
        (index, _) = closest_vertex(
            self.global_waypoints,
            np.array([self.position[:2]]),
        )

        num_waypoint_ahead = max(5, int(LOCAL_PLANNER_NUM_GLOBAL_WAYPOINTS_AHEAD * abs(self.speed[0] / LOCAL_PLANNER_DEFAULT_SPEED)))
        self.global_waypoints = self.global_waypoints[
            index : min(index + num_waypoint_ahead, len(self.global_waypoints))
        ]

        # return obstacles that are near global waypoints
        obstacle_list = get_obstacle_list(self.obstacles, self.global_waypoints[:, :2], OBSTACLE_PADDING)

        if len(obstacle_list) == 0:
            # Do not use Hybrid A* if there are no obstacles.
            speeds = self._adjust_speeds(self.global_waypoints, self.speed[0], self.target_speed)
            return self.global_waypoints, speeds, True

        # Hybrid a* does not take into account the driveable region.
        # It constructs search space as a top down, minimum bounding
        # rectangle with padding in each dimension.

        initial_conditions = self._compute_initial_conditions(obstacle_list)
        path_xs, path_ys, path_yaws, success = hybrid_astar_wrapper.apply_hybrid_astar(
            initial_conditions, self._hyperparameters
        )
        if DEBUG:
            print(f"[Hybrid A*] initial_conditions: {initial_conditions}, success: {success}")

        if not success:
            if DEBUG:
                print("[Hybrid A*] Compute local waypoints failed!")
            # TODO: if failed, should replan the path
            target_speeds = np.array([self.target_speed] * len(self.global_waypoints))
            return self.global_waypoints, target_speeds, False
        else:
            local_waypoints = np.array([path_xs, path_ys]).T
            target_speeds = self._adjust_speeds(local_waypoints, self.speed[0], self.target_speed)
            return local_waypoints, target_speeds, True

    def _compute_initial_conditions(self, obstacle_list):
        [x, y, _, qx, qy, qz, qw] = self.position
        [_, _, yaw] = R.from_quat([qx, qy, qz, qw]).as_euler(
            "xyz", degrees=False
        )

        # start point
        start = np.array([x, y, yaw])

        # end point
        if len(self.global_waypoints) == 0:
            # If no more waypoints left. Then current location is our end point.
            end = np.array([x, y, yaw])
        else:
            num_waypoint_ahead = max(5, int(LOCAL_PLANNER_NUM_GLOBAL_WAYPOINTS_AHEAD * abs(self.speed[0] / LOCAL_PLANNER_DEFAULT_SPEED)))
            end_index = min(num_waypoint_ahead, len(self.global_waypoints) - 1)
            [end_x, end_y] = self.global_waypoints[end_index]
            if len(self.global_waypoints) == 1:
                # If only one waypoint left
                end_yaw = yaw
            else:
                dx = self.global_waypoints[end_index, 0] - self.global_waypoints[end_index - 1, 0]
                dy = self.global_waypoints[end_index, 1] - self.global_waypoints[end_index - 1, 1]
                end_yaw = np.arctan2(dy, dx)
            end = np.array([end_x, end_y, end_yaw])

        initial_conditions = {
            "start": start,
            "end": end,
            "obs": obstacle_list,
        }
        
        return initial_conditions

    def _downsample_path(self, path_xs, path_ys, interval=3):
        # select every interval-th point
        if len(path_xs) < interval:
            return path_xs, path_ys

        path_xs_downsampled = path_xs[::interval]
        path_ys_downsampled = path_ys[::interval]

        return path_xs_downsampled, path_ys_downsampled

    def _smooth_path(self, path_x, path_y, num_points=50):
        tck = splrep(path_x, path_y, s=0)
        new_path_x = np.linspace(path_x[0], path_x[-1], num_points)
        new_path_y = splev(new_path_x, tck, der=0)

        return new_path_x, new_path_y

    def _adjust_speeds(self, waypoints, current_speed, target_speed):
        # number of waypoints
        num_waypoints = len(waypoints)

        if num_waypoints == 0:
            return np.array([])
        
        if num_waypoints == 1:
            return np.array([target_speed])

        # calculate curvature at each waypoint
        curvatures = self._calculate_curvatures(waypoints)

        # adjust speed
        speeds = np.ones(num_waypoints) * target_speed
        for i in range(num_waypoints):
            # adjust speed according to curvature
            # # option 1:
            # curvature_change = curvatures[i] - curvatures[i-1]
            # speed_change = LOCAL_PLANNER_MAX_LAT_ACCEL * curvature_change
            
            # acceleration = min(LOCAL_PLANNER_MAX_LAT_ACCEL, (target_speed - current_speed) / 1)
            # delta_v = acceleration * 1

            # current_speed = min(current_speed + delta_v, target_speed)
            # speeds[i] = current_speed + speed_change

            # # option 2:
            # speeds[i] = np.clip(speeds[i] + speed_change, 0, target_speed)
            
            # option 3:
            if curvatures[i] > 1e-8:  # avoid division by zero
                speed_based_on_curvature = np.sqrt(LOCAL_PLANNER_MAX_LAT_ACCEL / curvatures[i])
                speeds[i] = min(speeds[i], speed_based_on_curvature)

        return speeds

    def _calculate_curvatures(self, waypoints):
        # calculate curvature at each waypoint
        num_waypoints = len(waypoints)
        curvatures = np.zeros(num_waypoints)

        for i in range(1, num_waypoints-1):
            dx1 = waypoints[i][0] - waypoints[i-1][0]
            dy1 = waypoints[i][1] - waypoints[i-1][1]
            dx2 = waypoints[i+1][0] - waypoints[i][0]
            dy2 = waypoints[i+1][1] - waypoints[i][1]

            curvature = abs(dx1 * dy2 - dx2 * dy1) / np.sqrt((dx1**2 + dy1**2) * (dx2**2 + dy2**2) + 1e-8)
            curvatures[i] = curvature

        curvatures[0] = curvatures[1]
        curvatures[-1] = curvatures[-2]
        
        return curvatures